Repository navigation
Quantize models on __init__ / from_config - #46324
ArthurZucker wants to merge 2 commits into
Conversation
…ned`) Constructing a model directly (`cls(config)` / `AutoModel.from_config`) with a `quantization_config` set on the config currently ignores it and returns a full-precision model. This wires up on-the-fly quantization for that path, mirroring `from_pretrained`. `PreTrainedModel.post_init` now calls `_maybe_quantize_on_init`, which (only on the construction root, and only for real non-meta weights) builds the `HfQuantizer`, captures the random-init weights, swaps modules via `preprocess_model`, and routes the weights back through the exact same `convert_and_load_state_dict_in_model` loader used by `from_pretrained` so the result is identical to loading those weights quantized. Gated by a new opt-in `HfQuantizer.supports_quantize_on_init` flag (enabled for FineGrainedFP8); quantizers needing calibration stay opt-out. A meta-device guard keeps the `from_pretrained` path untouched.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
vasqu
left a comment
There was a problem hiding this comment.
Overal lgtm, have a few smaller comments but nothing major
| `from_pretrained`, so the resulting model is byte-for-byte equivalent to loading those weights quantized. | ||
|
|
||
| It is a no-op unless all of the following hold: | ||
| - `self` is the construction root (the outermost model, see `__init__`); |
There was a problem hiding this comment.
Nit: we could list with 1., 2., 3. and reference below for each condition where we opt out (return)
| return | ||
| hf_quantizer.validate_environment(device_map=None, weights_only=True) | ||
|
|
||
| dtype = self.dtype |
There was a problem hiding this comment.
Is self.dtype really safe? Was always wondering, maybe we should read from config if anything in these instances because it should default to fp32 here in any case, no?
| return | ||
|
|
||
| hf_quantizer = AutoHfQuantizer.from_config(quantization_config, pre_quantized=False) | ||
| if not getattr(hf_quantizer, "supports_quantize_on_init", False): |
There was a problem hiding this comment.
I feel like here we could add a warning - if users pass it on init, they would expect the conversion but here it would "fail" silently
PR is not fully ready yet, will recheck later as well
…tion Drop the `supports_quantize_on_init` opt-in flag: all non-calibration quantizers can quantize freshly-initialized weights, so reuse the existing `requires_calibration` to gate instead. `_maybe_quantize_on_init` moves out of `PreTrainedModel` into `quantizers_utils.maybe_quantize_on_init`. Resolve the quantizer class without instantiating (new `AutoHfQuantizer.resolve`, factored out of `from_config`) so we can check `requires_calibration` before any backend import. Calibration-only methods (AWQ, GPTQ, ...) now emit a warning instead of silently returning a full-precision model. Read the dtype off the freshly-initialized parameters with a default-dtype fallback rather than `self.dtype`, which raises StopIteration on a param-less model; the parameters are the source of truth that matches the captured state_dict.
|
[For maintainers] Suggested jobs to run (before merge) run-slow: finegrained_fp8 |
Honors
config.quantization_configwhen building viacls(config)/AutoModel.from_config, mirroringfrom_pretrained's on-the-fly quantization. Opt-in per quantizer (supports_quantize_on_init), FP8 enabled; meta guard leavesfrom_pretraineduntouched.No open PR covers this. Tested: CPU regression + fp8 GPU-required parity locally; GPU path via new
FP8QuantizeOnInitTest.This IS a draft but it should help testing quantized models and etc